from langchain.chains.llm import LLMChain
from openai import OpenAI
from dotenv import load_dotenv
import streamlit as st
import mysql.connector
import os
import re
from langchain.chains import create_sql_query_chain
from langchain_openai import ChatOpenAI
from langchain.utilities import SQLDatabase
from langchain_core.prompts import FewShotPromptTemplate, PromptTemplate# 加载环境变量

# 加载环境变量
load_dotenv()

# 初始化 OpenAI 客户端（显式传递API_KEY）
client = OpenAI()

# 安全字符校验正则表达式
SAFE_TABLE_NAME_REGEX = r'^[a-zA-Z0-9_]+$'

def validate_table_name(table_name):
    """验证表名是否符合安全规范"""
    return re.match(SAFE_TABLE_NAME_REGEX, table_name) is not None

# 连接 MySQL 数据库（增加连接池配置）
def connect_to_mysql(host, user, password, database, port):
    try:
        conn = mysql.connector.connect(
            host=host,
            user=user,
            password=password,
            database=database,
            port=port,
            charset='utf8mb4',
            pool_size=5,  # 增加连接池
            pool_name="sql_gen_pool"
        )
        return conn
    except mysql.connector.Error as err:
        st.error(f"🔌 数据库连接失败: {err}")
        return None

# 获取所有表的结构（增加表名校验）
def get_table_schemas(conn):
    try:
        cursor = conn.cursor()
        cursor.execute("SHOW TABLES")
        tables = cursor.fetchall()
        schemas = {}

        for table in tables:
            table_name = table[0]
            if not validate_table_name(table_name):
                st.warning(f"跳过非法表名: {table_name}")
                continue

            try:
                cursor.execute(f"DESCRIBE `{table_name}`")  # 使用反引号包裹表名
                schema = cursor.fetchall()
                schemas[table_name] = schema
            except mysql.connector.Error as err:
                st.error(f"获取表 {table_name} 结构失败: {err}")

        cursor.close()
        return schemas
    except mysql.connector.Error as err:
        st.error(f"📊 获取表结构失败: {err}")
        return None

# 执行 SQL 查询（增加查询超时限制）
def execute_sql_query(conn, sql_query):
    try:
        cursor = conn.cursor()
        cursor.execute(sql_query)
        result = cursor.fetchall()
        columns = [desc[0] for desc in cursor.description]
        cursor.close()
        return columns, result, None
    except mysql.connector.Error as err:
        return [], [], str(err)
    except Exception as err:
        return [], [], f"未知错误: {str(err)}"

# 检测语言（增加缓存机制）
@st.cache_data(ttl=3600)
def detect_language(text):
    try:
        response = client.chat.completions.create(
            model="gpt-3.5-turbo",  # 使用正确的模型名称
            messages=[{
                "role": "user",
                "content": f"帮我判断文本使用什么语言，只返回语言代码（zh/en/es等）：\n文本：{text}"
            }],
            max_tokens=10
        )
        return response.choices[0].message.content.strip().lower()
    except Exception as err:
        st.error(f"🌐 语言检测失败: {err}")
        return None

# 使用 LangChain 的 create_sql_query_chain 重构代码
def main():
    st.set_page_config(page_title="智能SQL生成器", layout="wide")

    with st.sidebar.expander("⚙️ 数据库配置", expanded=True):
        db_host = st.text_input("主机地址", value=os.getenv("MYSQL_HOST", "localhost"))
        db_user = st.text_input("用户名", value=os.getenv("MYSQL_USER", "root"))
        db_password = st.text_input("密码", type="password", value=os.getenv("MYSQL_PASSWORD", ""))
        db_name = st.text_input("数据库名", value=os.getenv("MYSQL_DATABASE", "testdb"))
        db_port = st.number_input("端口", value=int(os.getenv("MYSQL_PORT", 3306)))

    # 使用columns布局
    col1, col2 = st.columns([3, 2])

    with col1:
        st.header("📝 需求输入")
        user_input = st.text_area(
            "请输入您的数据需求（例如：查询北京地区最近三个月的订单）",
            height=150,
            help="支持自然语言描述，建议包含：\n- 需要查询的字段\n- 过滤条件\n- 排序方式"
        )

    if st.button("🚀 生成并执行", use_container_width=True):
        if not user_input.strip():
            st.warning("请输入有效需求")
            return

        conn = None
        try:
            # 连接数据库
            with st.spinner("🔗 连接数据库中..."):
                conn = connect_to_mysql(db_host, db_user, db_password, db_name, db_port)
                if not conn:
                    return

            # 获取表结构
            with st.spinner("📂 获取数据库元数据..."):
                table_schemas = get_table_schemas(conn)
                if not table_schemas:
                    return

            # 语言检测
            lang_code = detect_language(user_input) or "zh"

            db = SQLDatabase.from_uri(f"mysql+pymysql://{db_user}:{db_password}@{db_host}:{db_port}/{db_name}")

            llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)

            # examples = [
            #     {"input": "List all artists.", "query": "SELECT * FROM Artist;"},
            #     {
            #         "input": "Find all albums for the artist 'AC/DC'.",
            #         "query": "SELECT * FROM Album WHERE ArtistId = (SELECT ArtistId FROM Artist WHERE Name = 'AC/DC');",
            #     },
            #     {
            #         "input": "List all tracks in the 'Rock' genre.",
            #         "query": "SELECT * FROM Track WHERE GenreId = (SELECT GenreId FROM Genre WHERE Name = 'Rock');",
            #     },
            #     {
            #         "input": "Find the total duration of all tracks.",
            #         "query": "SELECT SUM(Milliseconds) FROM Track;",
            #     },
            #     {
            #         "input": "List all customers from Canada.",
            #         "query": "SELECT * FROM Customer WHERE Country = 'Canada';",
            #     },
            #     {
            #         "input": "How many tracks are there in the album with ID 5?",
            #         "query": "SELECT COUNT(*) FROM Track WHERE AlbumId = 5;",
            #     },
            #     {
            #         "input": "Find the total number of invoices.",
            #         "query": "SELECT COUNT(*) FROM Invoice;",
            #     },
            #     {
            #         "input": "List all tracks that are longer than 5 minutes.",
            #         "query": "SELECT * FROM Track WHERE Milliseconds > 300000;",
            #     },
            #     {
            #         "input": "Who are the top 5 customers by total purchase?",
            #         "query": "SELECT CustomerId, SUM(Total) AS TotalPurchase FROM Invoice GROUP BY CustomerId ORDER BY TotalPurchase DESC LIMIT 5;",
            #     },
            #     {
            #         "input": "Which albums are from the year 2000?",
            #         "query": "SELECT * FROM Album WHERE strftime('%Y', ReleaseDate) = '2000';",
            #     },
            #     {
            #         "input": "How many employees are there",
            #         "query": 'SELECT COUNT(*) FROM "Employee"',
            #     },
            # ]
            #
            # # 定义示例提示模板
            # example_prompt = PromptTemplate.from_template("User input: {input}\nSQL query: {query}")
            #
            # # 定义 FewShotPromptTemplate
            # prompt = FewShotPromptTemplate(
            #     examples=examples,
            #     example_prompt=example_prompt,
            #     prefix="""
            #     你是一个SQL专家。请根据以下问题生成一个SQL查询。
            #     问题: {input}  # 修改了这里的变量引用
            #     可用的表信息: {table_info}
            #     请确保查询只使用上述表，并且最多返回 {top_k} 条结果。
            #     """,
            #     suffix="User input: {input}\nSQL query: ",
            #     input_variables=["input", "top_k", "table_info"],
            # )
            # 有提示词
            # chain = create_sql_query_chain(llm, db, prompt)

            # 无提示词
            chain = create_sql_query_chain(llm, db)
            # 生成 SQL 查询
            with st.spinner("🧠 正在分析需求并生成SQL..."):
                sql_query = chain.invoke({"question": user_input})
                # 基础语法校验
                if not sql_query.lower().startswith(("select", "show", "describe")):
                    st.error(f"🤖 SQL生成失败: {sql_query}")
                    raise ValueError("仅支持查询类SQL（SELECT/SHOW/DESCRIBE）")

            # 显示生成的SQL
            with col2:
                st.header("📄 生成SQL")
                st.code(sql_query, language="sql")

            # 执行与重试逻辑
            max_retries = 3
            history = []
            for attempt in range(max_retries):
                with st.spinner(f"⚡ 执行查询（尝试 {attempt + 1}/{max_retries}）..."):
                    columns, result, error = execute_sql_query(conn, sql_query)

                if error:
                    st.error(f"❌ 执行错误：{error}")
                    with st.expander("查看错误详情"):
                        st.code(f"错误SQL:\n{sql_query}\n\n错误信息:\n{error}")

                    # 使用 LangChain 自动纠错
                    fix_prompt_template = PromptTemplate(
                        input_variables=["table_schemas", "user_input", "error_message", "history"],
                        template=(
                            "表结构：\n{table_schemas}\n需求：{user_input}\n"
                            "错误：{error_message}\n历史尝试：{history}\n"
                            "请：\n1. 分析错误原因\n2. 生成修正SQL\n3. 添加注释说明修改点"
                            "要求：\n1. 只返回SQL语句字符串"
                        )
                    )
                    fix_chain = LLMChain(llm=llm, prompt=fix_prompt_template, output_key="fixed_sql")

                    with st.spinner("🤖 正在尝试自动修复..."):
                        new_sql = fix_chain.run(table_schemas=table_schemas, user_input=user_input, error_message=error, history=history)
                        if new_sql and new_sql != sql_query:
                            history.append(sql_query)
                            sql_query = new_sql
                            st.code(f"修正SQL:\n{new_sql}", language="sql")
                            continue
                        else:
                            st.error("自动修复失败")
                            break
                else:
                    # 显示结果
                    with col2:
                        st.header("📊 查询结果")
                        if result:
                            st.info(f"共找到 {len(result)} 条记录")
                            st.dataframe(
                                data=result,
                                column_config={col: {"width": 200} for col in columns},
                                height=400,
                                use_container_width=True
                            )
                        else:
                            st.info("未找到匹配数据")

                    # 生成分析报告
                    with st.spinner("📝 生成分析报告..."):

                        analysis_prompt_template = PromptTemplate(
                            input_variables=["user_input", "result", "lang_code"],
                            template=(
                                f"用户的需求是：{user_input}\n\n"
                                f"查询结果是：\n{result}\n\n"
                                f"请以私人助理的语气，用自然、口语化的方式回复用户，简要总结查询结果并提供建议。"
                                f"请使用语言代码为 '{lang_code}' 的语言进行回复。"
                                f"如果查询结果为空，请高情商地提示用户重新提问。"
                            )
                        )

                        analysis_chain = LLMChain(llm=llm, prompt=analysis_prompt_template, output_key="analysis")
                        analysis = analysis_chain.run(user_input=user_input, query_result=result, language_code=lang_code)
                        st.header(" 查询结果")
                        st.success(analysis)
                    break
            else:
                st.error(f"经过 {max_retries} 次尝试仍无法执行成功")

        finally:
            if conn:
                try:
                    conn.close()
                except mysql.connector.errors.OperationalError as e:
                    st.warning(f"⚠️ 关闭数据库连接时发生错误: {e}")
                st.cache_data.clear()  # 清除缓存

if __name__ == "__main__":
    main()